from fl.util.chain import BlockchainManager

class Listener(BlockchainManager):
    def __init__(self, config):
        super(Listener, self).__init__(config)
        tx = self.contract.functions.startListen().transact({
            'from': self.account,
            'nonce': self.w3.eth.getTransactionCount(self.account)})
        self.w3.eth.wait_for_transaction_receipt(tx)

        self.queue = []

    def listen(self):

        # 局部模型上传完成监听
        res = self.contract.functions.listenFinishedLocalModelCollectState().call()
        if res == 1:
            self.queue.append({'name':'LocalModelUpdatedFinished', 'args':[]})

        # 节点vrf注册完成监听
        res = self.contract.functions.listenFinishedRegistryState().call()
        if res == 1:
            self.queue.append({'name':'StartSelectAggregators', 'args':[]})

        # 聚合节点选举完成监听
        res = self.contract.functions.listenSelectAggregatorsState().call()
        if res[0] == 1:
            self.queue.append({'name':'ConfirmCandidateAggregators', 'args':res[1]})

        # 候选模型上传完成监听
        res = self.contract.functions.listenFinishedCandidateState().call()
        if res == 1:
            self.queue.append({'name':'StartVote', 'args':[]})

        # 投票完成监听
        res = self.contract.functions.listenFinishedVoteState().call()
        if res == 1:
            self.queue.append({'name':'VoteFinished', 'args':[]})

        # 最终聚合节点确认完成监听
        res = self.contract.functions.listenFinalAggregatorState().call()
        if res[0] == 1:
            self.queue.append({'name':'AggregatorSelected', 'args':res[1]})

        # 全局模型上传完成监听
        res = self.contract.functions.listenGlobalModelUpdatesState().call()
        if res == 1:
            self.queue.append({'name':'GlobalModelUpdated', 'args':res})

        # 完整联邦学习训练完成监听
        res = self.contract.functions.listenFinishedState().call()
        if res == 1:
            tx = self.contract.functions.initListener().transact({
                'from': self.account,
                'nonce': self.w3.eth.getTransactionCount(self.account)
                })
            self.w3.eth.wait_for_transaction_receipt(tx)

            self.queue.append({'name':'Finished', 'args':[]})

    def get_events(self):
        return self.queue





